%% Visualize photon shot noise for metameric spectra
% Compare the photon shot noise produced by a true spectral signal, and a
% metamer. Also compare reconstruction of metamers as a function of image noise.

% Bernard Llanos
% Supervised by Dr. Y.H. Yang
% University of Alberta, Department of Computing Science
% File created June 17, 2019

% Presently, this script does not automatically save any output
parameters_list = {};

%% Input data and parameters

% Colour space conversion data
color_map_filename = fullfile('.', 'demo_data', 'multispectral_images', 'sensor.mat');

% Sample spectral reflectances
reflectances_filename = fullfile('.', 'demo_data', 'spectral_data', 'spectra_averaged.csv');

% Which samples to use from the file of reflectances
reflectance_columns = 13:36; % All possible samples
% Chosen index within the samples
selected_reflectance_index = 16;

% CIE D-illuminant
illuminant_filename = '${FILEPATH}';
illuminant_temperature = 6504; % From https://en.wikipedia.org/wiki/Standard_illuminant#Illuminant_series_D

% Spectral smoothness weight
regularization_weight = 1e-3;

% Image noise fraction
noise_fraction = 0.1;

% Image patch size
border = 10; % Border region to be discarded
image_sampling = [32, 32];
image_sampling_full = image_sampling + border * 2;

% Number of replicates to test effect of noise on ADMM spectral reconstruction
n_samples_reconstruction_admm = 10;

% Number of replicates to test effect of noise on smoothness-based spectral reconstruction
n_samples_reconstruction_smoothness = 100;

% Number of replicates to test effect of photon shot noise no colour
n_samples_projection = 1000;

% Maximum number of colour samples to plot
n_samples_plot = 1000;

% Quantiles to use as bounds on variable results
quantiles = [0.05, 0.95];

% Parameters which do not usually need to be changed
run('SetFixedParameters.m')

%% Load calibration data

[sensor_map, ~, bands_color] = loadColorMap(color_map_filename, false);

%% Load reflectance data

sample_table = readtable(reflectances_filename);
variable_names = sample_table.Properties.VariableNames;
bands_measured = sample_table.(variable_names{1});
reflectances = sample_table{:, reflectance_columns};

disp('Name of selected reflectance:')
disp(variable_names{reflectance_columns(selected_reflectance_index)});
fprintf('\n');

%% Load the illuminant and convert reflectance to radiance

illuminant_data = csvread(illuminant_filename);
bands_illuminant = illuminant_data(:, 1);
S_illuminant = illuminant_data(:, 2:end);
spd_illuminant = ciedIlluminant(...
    illuminant_temperature, bands_illuminant, S_illuminant, bands_illuminant...
);

[bands_radiance, radiances] = reflectanceToRadiance(...
    bands_illuminant, spd_illuminant,...
    bands_measured, reflectances...
);
n_bands_radiance = length(bands_radiance);

%% Convert to colour

color_weights_options = findSamplingOptions;
color_weights_options.interpolant = findSamplingOptions.interpolant_ref;
color_weights = colorWeights(...
    sensor_map, bands_color, bands_radiance, color_weights_options...
);

c = channelConversion(radiances.', color_weights, 2);

%% Find a metamer under ideal conditions

radiance_gt = radiances(:, selected_reflectance_index);
c_gt = c(selected_reflectance_index, :);
radiance_metamer = smoothMetamer(c_gt, color_weights, regularization_weight, 0, true);
c_metamer = channelConversion(radiance_metamer, color_weights, 1);

% Sanity check that colours match
disp('True colour:')
disp(c_gt)
disp('Metamer colour:')
disp(c_metamer)

gt_plot_color = [0, 0, 0; 0.5, 0.5, 0.5];
metamer_plot_colors = [0, 0, 1; 0.5, 0.5, 1];
admm_plot_colors = [1, 0, 0; 1, 0.5, 0.5];

figure;
hold on
plot(...
    bands_radiance, radiance_gt,...
    'Color', gt_plot_color(1, :), 'LineWidth', 2, 'Marker', 'none'...
);
plot(...
    bands_radiance, radiance_metamer,...
    'Color', metamer_plot_colors(1, :), 'LineWidth', 2, 'Marker', 'none', 'LineStyle', ':'...
);
hold off
xlabel('Wavelength [nm]');
ylabel('Radiance');
title('Spectral radiances, without raw image noise');

% ADMM image estimation
I_rgb_gt = repmat(reshape(c_gt, 1, 1, []), image_sampling_full(1), image_sampling_full(2), 1);
I_raw_noNoise_gt = mosaic(I_rgb_gt, bayer_pattern);

solvePatchesSpectralOptions.patch_options.patch_size = image_sampling;
solvePatchesSpectralOptions.patch_options.padding = border;
solvePatchesSpectralOptions.patch_options.target_patch = repmat(border + 1, 1, 2);

[ ~, I ] = solvePatchesSpectral(...
    [], I_raw_noNoise_gt, bayer_pattern, [], sensor_map, bands_color,...
    solvePatchesSpectralOptions.sampling_options,...
    solvePatchesSpectralOptions.admm_options,...
    solvePatchesSpectralOptions.reg_options,...
    solvePatchesSpectralOptions.patch_options...
);
[~, spectral_weights, bands] = findSampling(...
  sensor_map, bands_color, bands_radiance, findSamplingOptions...
);
I_spectral = channelConversion(I, spectral_weights);
I_2D = reshape(I_spectral, [], n_bands_radiance);
I_bounds = quantile(I_2D, quantiles, 1);
I_mean = mean(I_2D, 1);
radiance_admm = I_mean;
c_admm = channelConversion(I_mean, color_weights, 2);
disp('ADMM metamer colour:')
disp(c_admm)

hold on
plot(...
    bands_radiance, I_bounds(1, :),...
    'Color', admm_plot_colors(2, :), 'LineWidth', 2, 'Marker', 'none', 'LineStyle', '--'...
);
plot(...
    bands_radiance, I_mean,...
    'Color', admm_plot_colors(1, :), 'LineWidth', 2, 'Marker', 'none', 'LineStyle', '--'...
);
plot(...
    bands_radiance, I_bounds(2, :),...
    'Color', admm_plot_colors(2, :), 'LineWidth', 2, 'Marker', 'none', 'LineStyle', '--'...
);
hold off

legend(...
    'True', 'Spectral smoothness metamer',...
    sprintf('Ours, %g quantile', quantiles(1)),...
    'Ours, mean',...
    sprintf('Ours, %g quantile', quantiles(2))...
);

%% Find distribution of metamers with and without noise

snr = noiseFractionToSNR(noise_fraction);

n_px = prod(image_sampling);
I_spectral_samples = zeros(n_px * n_samples_reconstruction_admm, n_bands_radiance);
for i = 1:n_samples_reconstruction_admm
    I_raw_gt = addNoise(I_raw_noNoise_gt, snr); 
    [ ~, I ] = solvePatchesSpectral(...
        [], I_raw_gt, bayer_pattern, [], sensor_map, bands_color,...
        solvePatchesSpectralOptions.sampling_options,...
        solvePatchesSpectralOptions.admm_options,...
        solvePatchesSpectralOptions.reg_options,...
        solvePatchesSpectralOptions.patch_options...
    );
    I_spectral = channelConversion(I, spectral_weights);
    I_spectral_samples((n_px * (i - 1) + 1):(n_px * i), :) = reshape(I_spectral, [], n_bands_radiance);
end
I_bounds = quantile(I_spectral_samples, quantiles, 1);
I_mean = mean(I_spectral_samples, 1);

radiance_metamer_samples = zeros(n_bands_radiance, n_samples_reconstruction_smoothness);
c_gt_samples = zeros(n_samples_reconstruction_smoothness, length(c_gt));
for i = 1:n_samples_reconstruction_smoothness
    I_raw_gt = addNoise(I_raw_noNoise_gt, snr);
    c_gt_samples(i, :) = addNoise(c_gt, snr);
    
    radiance_metamer_samples(:, i) = smoothMetamer(c_gt_samples(i, :), color_weights, regularization_weight);    
end
radiance_metamer_bounds = quantile(radiance_metamer_samples, quantiles, 2);
radiance_metamer_mean = mean(radiance_metamer_samples, 2);

% Plot spectral variability
figure;
hold on
plot(...
    bands_radiance, radiance_gt,...
    'Color', gt_plot_color(1, :), 'LineWidth', 2, 'Marker', 'none'...
);
plot(...
    bands_radiance, radiance_metamer_bounds(:, 1),...
    'Color', metamer_plot_colors(2, :), 'LineWidth', 2, 'Marker', 'none', 'LineStyle', ':'...
);
plot(...
    bands_radiance, radiance_metamer_mean,...
    'Color', metamer_plot_colors(1, :), 'LineWidth', 2, 'Marker', 'none', 'LineStyle', ':'...
);
plot(...
    bands_radiance, radiance_metamer_bounds(:, 2),...
    'Color', metamer_plot_colors(2, :), 'LineWidth', 2, 'Marker', 'none', 'LineStyle', ':'...
);

plot(...
    bands_radiance, I_bounds(1, :),...
    'Color', admm_plot_colors(2, :), 'LineWidth', 2, 'Marker', 'none', 'LineStyle', '--'...
);
plot(...
    bands_radiance, I_mean,...
    'Color', admm_plot_colors(1, :), 'LineWidth', 2, 'Marker', 'none', 'LineStyle', '--'...
);
plot(...
    bands_radiance, I_bounds(2, :),...
    'Color', admm_plot_colors(2, :), 'LineWidth', 2, 'Marker', 'none', 'LineStyle', '--'...
);
hold off

xlabel('Wavelength [nm]');
ylabel('Radiance');
title(sprintf('Spectral radiances, with %g raw image noise fraction', noise_fraction));

legend(...
    'True',...
    sprintf('Spectral smoothness metamer, %g quantile', quantiles(1)),...
    'Spectral smoothness metamer, mean',...
    sprintf('Spectral smoothness metamer, %g quantile', quantiles(2)),...
    sprintf('Ours, %g quantile', quantiles(1)),...
    'Ours, mean',...
    sprintf('Ours, %g quantile', quantiles(2))...
);

% Plot colour variability
plot_sample_indices = randperm(size(I_spectral_samples, 1), min(n_samples_plot, size(I_spectral_samples, 1)));
c_admm_samples = channelConversion(I_spectral_samples(plot_sample_indices, :), color_weights, 2);
plot_sample_indices = randperm(size(radiance_metamer_samples, 2), min(n_samples_plot, size(radiance_metamer_samples, 2)));
c_metamer_samples = channelConversion(radiance_metamer_samples(:, plot_sample_indices).', color_weights, 2);
plot_sample_indices = randperm(size(c_gt_samples, 1), min(n_samples_plot, size(c_gt_samples, 1)));
c_gt_samples_plot = c_gt_samples(plot_sample_indices, :);

figure;
hold on
scatter3(c_admm_samples(:, 1), c_admm_samples(:, 2), c_admm_samples(:, 3), [], admm_plot_colors(1, :), 'filled')
scatter3(c_metamer_samples(:, 1), c_metamer_samples(:, 2), c_metamer_samples(:, 3), [], metamer_plot_colors(1, :), 'filled')
scatter3(c_gt_samples_plot(:, 1), c_gt_samples_plot(:, 2), c_gt_samples_plot(:, 3), [], gt_plot_color(1, :), 'filled')
hold off

xlabel('Red');
ylabel('Green');
zlabel('Blue');
title(sprintf('Reconstructed colour variability, with %g raw image noise fraction', noise_fraction));

legend(...
    'Spectral smoothness metamer',...
    'Ours',...
    'True'...
);

%% Find distribution of RGB noise with noise in metamers

shot_samples_gt = shotNoise(reshape(radiance_gt, 1, 1, []), n_samples_projection);
shot_samples_gt = shiftdim(shot_samples_gt).';
radiance_gt_bounds = quantile(shot_samples_gt, quantiles, 1);

shot_samples_metamer = shotNoise(reshape(radiance_metamer, 1, 1, []), n_samples_projection);
shot_samples_metamer = shiftdim(shot_samples_metamer).';
radiance_metamer_bounds = quantile(shot_samples_metamer, quantiles, 1);

shot_samples_admm = shotNoise(reshape(radiance_admm, 1, 1, []), n_samples_projection);
shot_samples_admm = shiftdim(shot_samples_admm).';
I_bounds = quantile(shot_samples_admm, quantiles, 1);

% Plot spectral variability
figure;
hold on
plot(...
    bands_radiance, radiance_gt_bounds(1, :),...
    'Color', gt_plot_color(2, :), 'LineWidth', 2, 'Marker', 'none'...
);
plot(...
    bands_radiance, radiance_gt,...
    'Color', gt_plot_color(1, :), 'LineWidth', 2, 'Marker', 'none'...
);
plot(...
    bands_radiance, radiance_gt_bounds(2, :),...
    'Color', gt_plot_color(2, :), 'LineWidth', 2, 'Marker', 'none'...
);
plot(...
    bands_radiance, radiance_metamer_bounds(1, :),...
    'Color', metamer_plot_colors(2, :), 'LineWidth', 2, 'Marker', 'none', 'LineStyle', ':'...
);
plot(...
    bands_radiance, radiance_metamer,...
    'Color', metamer_plot_colors(1, :), 'LineWidth', 2, 'Marker', 'none', 'LineStyle', ':'...
);
plot(...
    bands_radiance, radiance_metamer_bounds(2, :),...
    'Color', metamer_plot_colors(2, :), 'LineWidth', 2, 'Marker', 'none', 'LineStyle', ':'...
);

plot(...
    bands_radiance, I_bounds(1, :),...
    'Color', admm_plot_colors(2, :), 'LineWidth', 2, 'Marker', 'none', 'LineStyle', '--'...
);
plot(...
    bands_radiance, radiance_admm,...
    'Color', admm_plot_colors(1, :), 'LineWidth', 2, 'Marker', 'none', 'LineStyle', '--'...
);
plot(...
    bands_radiance, I_bounds(2, :),...
    'Color', admm_plot_colors(2, :), 'LineWidth', 2, 'Marker', 'none', 'LineStyle', '--'...
);
hold off

xlabel('Wavelength [nm]');
ylabel('Radiance');
title('Spectral radiances, subject to Poisson shot noise');

legend(...
    sprintf('True, %g quantile', quantiles(1)),...
    'True',...
    sprintf('True, %g quantile', quantiles(2)),...
    sprintf('Spectral smoothness metamer, %g quantile', quantiles(1)),...
    'Spectral smoothness metamer',...
    sprintf('Spectral smoothness metamer, %g quantile', quantiles(2)),...
    sprintf('Ours, %g quantile', quantiles(1)),...
    'Ours',...
    sprintf('Ours, %g quantile', quantiles(2))...
);

% Plot colour variability
plot_sample_indices = randperm(size(shot_samples_admm, 1), min(n_samples_plot, size(shot_samples_admm, 1)));
c_admm_samples = channelConversion(shot_samples_admm(plot_sample_indices, :), color_weights, 2);
plot_sample_indices = randperm(size(shot_samples_metamer, 1), min(n_samples_plot, size(shot_samples_metamer, 1)));
c_metamer_samples = channelConversion(shot_samples_metamer(plot_sample_indices, :), color_weights, 2);
plot_sample_indices = randperm(size(shot_samples_gt, 1), min(n_samples_plot, size(shot_samples_gt, 1)));
c_gt_samples_plot = channelConversion(shot_samples_gt(plot_sample_indices, :), color_weights, 2);

figure;
hold on
scatter3(c_admm_samples(:, 1), c_admm_samples(:, 2), c_admm_samples(:, 3), [], admm_plot_colors(1, :), 'filled')
scatter3(c_metamer_samples(:, 1), c_metamer_samples(:, 2), c_metamer_samples(:, 3), [], metamer_plot_colors(1, :), 'filled')
scatter3(c_gt_samples_plot(:, 1), c_gt_samples_plot(:, 2), c_gt_samples_plot(:, 3), [], gt_plot_color(1, :), 'filled')
hold off

xlabel('Red');
ylabel('Green');
zlabel('Blue');
title('Colour variability with Poisson shot noise');

legend(...
    'Spectral smoothness metamer',...
    'Ours',...
    'True'...
);